import numpy as np
import yaml, copy


class KptTransToCrop:
    """
    将keypoint_detector中还原后的坐标转换为图片切片后的对应坐标
    args:
        keys:
            格式:keys={
                "in":"your_in_keyword",
                "crop_targets":"crop_targets",
                "out":"your_save_keyword"
            }
    """

    def __init__(self, keys):
        self.keys = keys

    def __call__(self, data):
        # 获取经过KPT_detector得到的data
        kpt_results = copy.deepcopy(data[self.keys["in"]])
        # 获取切片data
        crop_targets_data = data[self.keys["crop_targets"]]
        results = []
        for idx, data_ in enumerate(crop_targets_data):
            if len(kpt_results) > 0:
                single_img_res = kpt_results[idx]
                batch_records = np.array(data_["boxes"])
                keypoint_vector, score_vector = translate_to_new_images(
                    single_img_res, batch_records
                )
                single_img_res["keypoint"] = keypoint_vector
                single_img_res["score"] = score_vector
            else:
                single_img_res = {}
            results.append(single_img_res)
        data[self.keys["out"]] = results

        return data


# 将原图关键点的对应坐标还原到切片对应坐标上
def translate_to_new_images(keypoint_result, batch_records):
    kpts = keypoint_result["keypoint"]
    scores = keypoint_result["score"]
    kpts[..., 0] -= batch_records[:, 0:1]
    kpts[..., 1] -= batch_records[:, 1:2]
    return kpts, scores
